import datetime as dt
import pandas as pd
from tqdm.auto import tqdm
from shares.api import TushareClient
from shares.cache import cache
import numpy as np
import time


class TushareClientPlus(TushareClient):

    def __init__(self, token=None):
        super(TushareClientPlus, self).__init__(token)


    def trade_days(self, start_date=None, end_date=None, count=1000) -> np.ndarray:
        """
        获取指定范围交易日
        :param start_date: 开始日期
        :param end_date: 结束日期
        :param count: 数量
        :return:
        """
        end_date = dt.datetime.now() if end_date is None else end_date
        key = "交易日_base"
        with cache:
            cached = cache.get(key)
        if cached is None:
            df_cal = self.trade_cal(end_date=(dt.datetime.now() + dt.timedelta(days=365)).strftime("%Y%m%d"),
                                    is_open='1').get().sort_values("cal_date").reset_index()
            days = np.array(pd.to_datetime(df_cal['cal_date']).tolist())
            with cache:
                cache.set(key, days, expire=86400, retry=True)
            cached = days
        if start_date is None:
            return cached[cached <= pd.to_datetime(end_date.strftime("%Y-%m-%d"))][-count:] # type: ignore
        else:
            return cached[(cached >= pd.to_datetime(start_date.strftime("%Y-%m-%d"))) & ( # type: ignore
                    cached <= pd.to_datetime(end_date.strftime("%Y-%m-%d")))] # type: ignore

    def is_open(self, date=None) -> bool:
        """
        给定日期是不是交易日
        :return:
        """
        date = dt.datetime.now() if date is None else date
        days = self.trade_days(end_date=date, count=1)
        return pd.to_datetime(date.strftime("%Y-%m-%d")) in days

    def last_closed_trade_day(self, end_date=None, end_minute=60 * 15 + 1) -> dt.datetime:
        """最后一个收盘的交易日"""
        end_date = dt.datetime.now() if end_date is None else end_date
        trade_days = self.trade_days(end_date=end_date, count=5)
        if (dt.datetime.now() - trade_days[-1]).days == 0 and end_date.hour * 60 + end_date.minute < end_minute:
            # 今天是交易日，沒有收盤
            return trade_days[-2]
        else:
            return trade_days[-1]

    @staticmethod
    def last_quarters(count=10, end_date=None):
        """
        最新报告期
        :param end_date:
        :param count:
        :return:
        """
        end_date = dt.datetime.now() if end_date is None else end_date
        return pd.date_range(end=end_date.date(), periods=count, freq="Q")

    @staticmethod
    def clean_cache(tags: str = "api,panel"):
        """
        清理指定类型的api缓存
        :param tags:清理指定tag的缓存
        :return:
        """
        with cache:
            tag_list = tags.split(",")
            for tag in tag_list:
                cache.evict(tag=tag)

    def _daily_bar_cache(self, date):
        """
        每日K线数据
        :param date:
        :return:
        """
        date = pd.to_datetime(date).strftime('%Y%m%d')
        key = "_daily_bar_cache-" + date
        with cache:
            cached = cache.get(key)
        if cached is None:
            df_price = self.daily(trade_date=date).get()
            if (df_price.shape[0]) > 0:
                df_price['vol'] = df_price['vol'].astype(np.float64) * 100
                df_price['amount'] = df_price['amount'].astype(np.float64) * 1000
                df_price['trade_date'] = pd.to_datetime(df_price['trade_date'])
                df_adj = self.adj_factor(trade_date=date).get()
                df_adj['trade_date'] = pd.to_datetime(df_adj['trade_date'])
                df_limit = self.stk_limit(trade_date=date).get()
                df_limit['trade_date'] = pd.to_datetime(df_limit['trade_date'])
                df_ret = df_price.merge(df_adj, left_on=['ts_code', 'trade_date'], right_on=['ts_code', 'trade_date'],
                                        how='left')
                df_ret1 = df_ret.merge(df_limit, left_on=['ts_code', 'trade_date'], right_on=['ts_code', 'trade_date'],
                                       how='left')
                df_ret2 = df_ret1[
                    ['ts_code', 'trade_date', 'open', 'high', 'low', 'close', 'vol', 'amount', 'pre_close', 'change',
                     'pct_chg', 'adj_factor', 'up_limit',
                     'down_limit']].dropna().copy()
                if df_ret2.shape[0] > 0:
                    with cache:
                        cache.set(key, df_ret2, retry=True, tag='panel')
                return df_ret2
            else:
                return pd.DataFrame()
        else:
            return cached

    def _daily_valuation_cache(self, date):
        """
        每日估值数据
        :param date:
        :return:
        """
        date = pd.to_datetime(date).strftime('%Y%m%d')
        key = "_daily_valuation_cache-" + date
        with cache:
            cached = cache.get(key)
        if cached is None:
            df = self.daily_basic(trade_date=date).get()
            df['total_mv'] = df['total_mv'].astype(np.float64) * 10000
            df['circ_mv'] = df['circ_mv'].astype(np.float64) * 10000
            if df.shape[0] > 0:
                with cache:
                    cache.set(key, df, retry=True, tag='panel')
            return df
        else:
            return cached

    def _daily_moneyflow_cache(self, date):
        """
        每日资金流
        :param date:
        :return:
        """
        date = pd.to_datetime(date).strftime('%Y%m%d')
        key = "_daily_moneyflow_cache-" + date
        with cache:
            cached = cache.get(key)
        if cached is None:
            df = self.moneyflow(trade_date=date).get()
            if df.shape[0] > 0:
                with cache:
                    cache.set(key, df, retry=True, tag='panel')
            return df
        else:
            return cached

    def _daily_margin_cache(self, date):
        """
        每日融资融券数据
        :param date:
        :return:
        """
        date = pd.to_datetime(date).strftime('%Y%m%d')
        key = "_daily_margin_cache-" + date
        with cache:
            cached = cache.get(key)
        if cached is None:
            df = self.margin_detail(trade_date=date).get()
            if df.shape[0] > 0:
                with cache:
                    cache.set(key, df, retry=True, tag='panel')
            time.sleep(0.12)
            return df
        else:
            return cached

    @staticmethod
    def _adj_fields(df_pivot: pd.DataFrame, arr_adj: pd.DataFrame) -> pd.DataFrame:
        """
        针对tushare 以千,万,手为单位的字段,转换成1为单位
        :param df_pivot:
        :param arr_adj:
        :return:
        """
        return pd.DataFrame(df_pivot.values * arr_adj, columns=df_pivot.columns,
                            index=df_pivot.index)

    def panel_value(self, end_date=None, count=255, use_cache=True):
        """
        估值数据面板
        :param end_date:
        :param count:
        :param use_cache:
        :return:
        """
        end_date = dt.datetime.now() if end_date is None else end_date
        last_day = self.last_closed_trade_day(end_date)
        trade_days = self.trade_days(end_date=last_day, count=count)
        dfs = []
        for date in tqdm(trade_days):
            dfs.append(self._daily_valuation_cache(date))
        df_all = pd.concat(dfs)
        df_all['trade_date'] = pd.to_datetime(df_all['trade_date'])
        df_turnover_rate = df_all.pivot(index='trade_date', columns='ts_code', values='turnover_rate')
        df_turnover_rate_f = df_all.pivot(index='trade_date', columns='ts_code', values='turnover_rate_f')
        df_volume_ratio = df_all.pivot(index='trade_date', columns='ts_code', values='volume_ratio')
        df_pe_ttm = df_all.pivot(index='trade_date', columns='ts_code', values='pe_ttm')
        df_pb = df_all.pivot(index='trade_date', columns='ts_code', values='pb')
        df_close = df_all.pivot(index='trade_date', columns='ts_code', values='close')
        df_total_mv = df_all.pivot(index='trade_date', columns='ts_code', values='total_mv')
        df_circ_mv = df_all.pivot(index='trade_date', columns='ts_code', values='circ_mv')
        ret = {"换手率": df_turnover_rate, "流通换手率": df_turnover_rate_f,
               '量比': df_volume_ratio, 'PE_TTM': df_pe_ttm, 'PB': df_pb,
               '总市值(亿)': df_total_mv, '流通市值(亿)': df_circ_mv,
               'close': df_close
               }
        return ret

    def panel_kline(self, end_date=None, count=255, use_cache=True):
        """
        股票K线Panel
        :param end_date:
        :param count:
        :param use_cache:
        :return:
        """
        end_date = dt.datetime.now() if end_date is None else end_date
        last_day = self.last_closed_trade_day(end_date)
        trade_days = self.trade_days(end_date=last_day, count=count + 1)
        data_missing = False
        with cache:
            df_old = cache.get("all_kline")
        old_days = [] if df_old is None else list(pd.to_datetime(df_old['trade_date'].unique())) # type: ignore
        need_update_days = set(trade_days.tolist()).difference(set(old_days))
        if df_old is None:
            dfs = []
        else:
            dfs = [df_old]
        for date in tqdm(need_update_days):
            df_day = self._daily_bar_cache(date)
            if df_day is not None and df_day.shape[0] > 0: # type: ignore
                dfs.append(df_day)
            else:
                data_missing = True
        df_all = pd.concat(dfs) # type: ignore
        types = {
            'open': 'float32',
            'high': 'float32',
            'low': 'float32',
            'close': 'float32',
            'vol': 'float32',
            'amount': 'float32',
            'adj_factor': 'float32',
            'up_limit': 'float32',
            'down_limit': 'float32',
            'ts_code': 'category'
        }
        df_all = df_all.drop_duplicates(subset=['ts_code', 'trade_date']).astype(types, errors='ignore')
        cache.set("all_kline", df_all, 2592000, tag='panel')
        df_all = df_all[(df_all['trade_date'] >= trade_days[0]) & (df_all['trade_date'] <= trade_days[-1])].copy()
        df_all_pivot = df_all.pivot(index='trade_date', columns='ts_code',
                                    values=['adj_factor', 'open', 'high', 'low', 'close', 'vol', 'amount',
                                            'up_limit', 'down_limit'])
        df_adj = df_all_pivot['adj_factor'].ffill().bfill().fillna(1).astype(
            np.float64, errors='ignore')
        arr_adj = df_adj.values / df_adj.iloc[-1].values
        df_open = df_all_pivot['open'].ffill().bfill().fillna(0)
        df_close = df_all_pivot['close'].ffill().bfill().fillna(0)
        df_high = df_all_pivot['high'].ffill().bfill().fillna(0)
        df_low = df_all_pivot['low'].ffill().bfill().fillna(0)
        df_volume = df_all_pivot['vol'].fillna(0)
        df_money = df_all_pivot['amount'].fillna(0)
        df_up_limit = df_all_pivot['up_limit'].fillna(10000000)
        df_down_limit = df_all_pivot['down_limit'].fillna(0)
        df_open = TushareClientPlus._adj_fields(df_open, arr_adj).iloc[1:]
        df_close = TushareClientPlus._adj_fields(df_close, arr_adj)
        df_pre_close = df_close.shift(1).iloc[1:]
        df_close = df_close.iloc[1:]
        df_high = TushareClientPlus._adj_fields(df_high, arr_adj).iloc[1:]
        df_low = TushareClientPlus._adj_fields(df_low, arr_adj).iloc[1:]
        df_high_limit = TushareClientPlus._adj_fields(df_up_limit,
                                                      arr_adj).iloc[1:]
        df_low_limit = TushareClientPlus._adj_fields(df_down_limit,
                                                     arr_adj).iloc[1:]
        df_volume = pd.DataFrame((df_volume.values / arr_adj), columns=df_volume.columns,
                                 index=df_volume.index)[1:]
        df_money = pd.DataFrame(df_money.values,
                                columns=df_money.columns, index=df_money.index).iloc[1:]
        ret = {"open": df_open, "high": df_high, 'low': df_low, 'close': df_close, 'vol': df_volume,
               'amount': df_money, 'pre_close': df_pre_close, 'high_limit': df_high_limit, 'low_limit': df_low_limit
               }
        return ret

    def panel_margin(self, end_date=None, count=255, use_cache=True):
        """
        融资融券面板
        :param end_date:
        :param count:
        :param use_cache:
        :return:
        """
        end_date = (dt.datetime.now() - dt.timedelta(days=1)) if end_date is None else end_date
        last_day = self.last_closed_trade_day(end_date)
        trade_days = self.trade_days(end_date=last_day, count=count)
        dfs = []
        for date in tqdm(trade_days):
            dfs.append(self._daily_margin_cache(date))
        df_all = pd.concat(dfs)
        df_all['trade_date'] = pd.to_datetime(df_all['trade_date'])
        df_rzye = df_all.pivot(index='trade_date', columns='ts_code', values='rzye').fillna(0)
        df_rqye = df_all.pivot(index='trade_date', columns='ts_code', values='rqye').fillna(0)
        df_rzmre = df_all.pivot(index='trade_date', columns='ts_code', values='rzmre').fillna(0)
        df_rqyl = df_all.pivot(index='trade_date', columns='ts_code', values='rqyl').fillna(0)
        df_rzche = df_all.pivot(index='trade_date', columns='ts_code', values='rzche').fillna(0)
        df_rqchl = df_all.pivot(index='trade_date', columns='ts_code', values='rqchl').fillna(0)
        df_rqmcl = df_all.pivot(index='trade_date', columns='ts_code', values='rqmcl').fillna(0)
        df_rzrqye = df_all.pivot(index='trade_date', columns='ts_code', values='rzrqye').fillna(0)
        ret = {"融资余额": df_rzye, "融券余额": df_rqye,
               '融资买入额': df_rzmre, '融券余量': df_rqyl, '融资偿还额': df_rzche,
               '融券偿还量': df_rqchl, '融券卖出量': df_rqmcl, '融资融券余额': df_rzrqye
               }
        return ret

    def panel_moneyflow(self, end_date=None, count=255, use_cache=True):
        """
        资金流面板数据
        :param end_date:
        :param count:
        :param use_cache:
        :return:
        """
        end_date = dt.datetime.now() if end_date is None else end_date
        last_day = self.last_closed_trade_day(end_date)
        trade_days = self.trade_days(end_date=last_day, count=count)
        dfs = []
        for date in tqdm(trade_days):
            dfs.append(self._daily_moneyflow_cache(date))
        df_all = pd.concat(dfs)
        df_all['trade_date'] = pd.to_datetime(df_all['trade_date'])
        buy_sm_vol = df_all.pivot(index='trade_date', columns='ts_code', values='buy_sm_vol').fillna(0) * 100
        buy_sm_amount = df_all.pivot(index='trade_date', columns='ts_code', values='buy_sm_amount').fillna(
            0) * 10000
        sell_sm_vol = df_all.pivot(index='trade_date', columns='ts_code', values='sell_sm_vol').fillna(0) * 100
        sell_sm_amount = df_all.pivot(index='trade_date', columns='ts_code', values='sell_sm_amount').fillna(
            0) * 10000
        buy_md_vol = df_all.pivot(index='trade_date', columns='ts_code', values='buy_md_vol').fillna(0) * 100
        buy_md_amount = df_all.pivot(index='trade_date', columns='ts_code', values='buy_md_amount').fillna(
            0) * 10000
        sell_md_vol = df_all.pivot(index='trade_date', columns='ts_code', values='sell_md_vol').fillna(0) * 100
        sell_md_amount = df_all.pivot(index='trade_date', columns='ts_code', values='sell_md_amount').fillna(
            0) * 10000
        buy_lg_vol = df_all.pivot(index='trade_date', columns='ts_code', values='buy_lg_vol').fillna(0) * 100
        buy_lg_amount = df_all.pivot(index='trade_date', columns='ts_code', values='buy_lg_amount').fillna(
            0) * 10000
        sell_lg_vol = df_all.pivot(index='trade_date', columns='ts_code', values='sell_lg_vol').fillna(0) * 100
        sell_lg_amount = df_all.pivot(index='trade_date', columns='ts_code', values='sell_lg_amount').fillna(
            0) * 10000
        buy_elg_vol = df_all.pivot(index='trade_date', columns='ts_code', values='buy_elg_vol').fillna(0) * 100
        buy_elg_amount = df_all.pivot(index='trade_date', columns='ts_code', values='buy_elg_amount').fillna(
            0) * 10000
        sell_elg_vol = df_all.pivot(index='trade_date', columns='ts_code', values='sell_elg_vol').fillna(0) * 100
        sell_elg_amount = df_all.pivot(index='trade_date', columns='ts_code', values='sell_elg_amount').fillna(
            0) * 10000
        net_mf_vol = df_all.pivot(index='trade_date', columns='ts_code', values='net_mf_vol').fillna(0) * 100
        net_mf_amount = df_all.pivot(index='trade_date', columns='ts_code', values='net_mf_amount').fillna(
            0) * 10000
        ret = {"小单买入量": buy_sm_vol, "小单买入金额": buy_sm_amount, '小单卖出量': sell_sm_vol,
               '小单卖出金额': sell_sm_amount,
               '中单买入量': buy_md_vol, '中单买入金额': buy_md_amount, '中单卖出量': sell_md_vol,
               '中单卖出金额': sell_md_amount,
               '大单买入量': buy_lg_vol, '大单买入金额': buy_lg_amount, '大单卖出量': sell_lg_vol,
               '大单卖出金额': sell_lg_amount,
               '特大单买入量': buy_elg_vol, '特大单买入金额': buy_elg_amount, '特大单卖出量': sell_elg_vol,
               '特大单卖出金额': sell_elg_amount,
               '净流入量': net_mf_vol, '净流入额': net_mf_amount
               }
        return ret

    def series_name(self, s):
        """
        根据code列返回name列
        :param s:
        :return:
        """
        name_dict = self.stock_basic().cache_or_get().set_index("ts_code")['name'].to_dict()
        keys = name_dict.keys()
        if "series" in str(type(s)):
            return s.apply(lambda cell: name_dict[cell] if cell in keys else cell)
        elif "index" in str(type(s)):
            s1 = pd.Series(index=s, data=s)
            return s1.apply(lambda cell: name_dict[cell] if cell in keys else cell)
